# $Id: CompanionData.py 76 2010-01-12 17:13:31Z dark.ryder $

# Copyright (c) 2009, Ben Blank
#
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# * Redistributions of source code must retain the above copyright
#   notice, this list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright
#   notice, this list of conditions and the following disclaimer in the
#   documentation and/or other materials provided with the distribution.
#
# * Neither the name of 535 Design nor the names of its contributors
#   may be used to endorse or promote products derived from this
#   software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED.	IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import lxml.etree
import time
import urllib.request

start = time.clock()

terrains = { "flying": {}, "ground": {}, "swimming": {} }
errors = []

mount_ids_ignore = [] # none currently

mount_ids_adapt_flying = [
	71342, # Big Love Rocket
	48025, # Headless Horseman's Mount
	72286, # Invincible
	54729, # Winged Steed of the Ebon Blade
]

mount_ids_adapt_ground = [
	58983, # Big Blizzard Bear
	71342, # Big Love Rocket
	48025, # Headless Horseman's Mount
	72286, # Invincible
	65917, # Magic Rooster
]

mount_ids_cosmetic = [
	30174, # Riding Turtle
]

ns = { "html": "http://www.w3.org/1999/xhtml" }

parser = lxml.etree.XMLParser(recover=True)

for url in ["http://www.wowhead.com/?spells=-5", "http://www.wowhead.com/?spells&filter=minle=2;me=21"]:
	try:
		with urllib.request.urlopen(url) as page:
			doc = lxml.etree.parse(page, parser=parser)

		for link in doc.xpath("//html:div[@class='listview-void']/html:a", namespaces=ns):
			path = link.get("href")

			if not path.startswith("/?spell="):
				continue

			speeds = []
			id = int(path[8:])
			name = link.text.strip()

			if id in mount_ids_ignore:
				continue

			if id in mount_ids_cosmetic:
				speeds.append(("ground", 0))

			if id in mount_ids_adapt_flying:
				speeds.append(("flying", 150))
				speeds.append(("flying", 280))

			if id in mount_ids_adapt_ground:
				speeds.append(("ground", 60))
				speeds.append(("ground", 100))

			if len(speeds) == 0:
				try:
					spell_url = "http://www.wowhead.com/?spell=" + str(id)

					with urllib.request.urlopen(spell_url) as spell_page:
						spell_doc = lxml.etree.parse(spell_page, parser=parser)

					for cell in spell_doc.xpath("//html:td[html:small and starts-with(.,'Apply Aura: Mod Speed')]", namespaces=ns):
						effect = cell.text.strip()[22:] # clip off the known part ("Apply Aura: ...")
						value = int(cell.xpath("./html:small/html:br", namespaces=ns)[0].tail[7:])

						if effect.startswith("Mounted Flight"):
							# flying mounts can only be used in flyable areas and so aren't considered any other kind of mount
							speeds = [ ("flying", value) ]
							break
						elif effect.startswith("Swim"):
							# the swimming mount aura appears to apply the same speed bonus when out of water
							speeds.append(("ground", value))
							speeds.append(("swimming", value))
						elif effect.startswith("Mounted"):
							speeds.append(("ground", value))
				except urllib.request.URLError as ex:
					errors.append("Could not retrieve {0} ({1})".format(spell_url, ex))
					continue
				except lxml.etree.LxmlError as ex:
					errors.append("Could not parse {0} ({1})".format(spell_url, ex))
					continue

			if len(speeds) == 0:
				errors.append("No listed speeds for http://www.wowhead.com/?spell={0:5} ({1})".format(id, name))

			for terrain, speed in speeds:
				if not speed in terrains[terrain]:
					terrains[terrain][speed] = {}

				terrains[terrain][speed][id] = name
	except urllib.request.URLError as ex:
		errors.append("Could not retrieve {0} ({1})".format(url, ex))
	except lxml.etree.LxmlError as ex:
		errors.append("Could not parse {0} ({1})".format(url, ex))

try:
	with open("CompanionData.lua", "wb") as lua:
		lua.write(b"""\
-- this file was automatically generated by CompanionData.py

local Doolittle = LibStub("AceAddon-3.0"):GetAddon("Doolittle")

Doolittle.CRITTER = {
	pools = {
		-- if other pet reagents are ever added, this will be replaced
		-- with a proper generator like the one used for mounts
		i17202 = { -- Snowball
			26533, -- Father Winter's Helper
			26045, -- Tiny Snowman
			26529, -- Winter Reindeer
			26541, -- Winter's Little Helper
		},
	},
}

""")

		for error in errors:
			lua.write("-- {0}\n".format(error).encode("utf-8"))

		lua.write(b"""\
Doolittle.MOUNT = {
	pools = {
		aq40 = {
			25953, -- Blue Qiraji Battle Tank
			26056, -- Green Qiraji Battle Tank
			26054, -- Red Qiraji Battle Tank
			26055, -- Yellow Qiraji Battle Tank
		},
""")

		for terrain, speeds in sorted(terrains.items(), key=lambda item: item[0]):
			lua.write("\t\t{0} = {{\n".format(terrain).encode("utf-8"))

			for speed, mounts in sorted(speeds.items(), key=lambda item: item[0]):
				lua.write("\t\t\t[{0}] = Pool{{\n".format(speed).encode("utf-8"))

				for id, name in sorted(mounts.items(), key=lambda item: (item[1], item[0])):
					lua.write("\t\t\t\t{0:>5}, -- {1}\n".format(id, name).encode("utf-8"))

				lua.write(b"\t\t\t},\n")

			lua.write(b"\t\t},\n")

		lua.write(b"\t},\n\tspeeds = {\n")

		for terrain, speeds in sorted(terrains.items(), key=lambda item: item[0]):
			lua.write("\t\t{0} = {{\n".format(terrain).encode("utf-8"))

			speeds = list(speeds.keys())
			speeds.sort()
			limit = len(speeds) - 2

			for i, speed in enumerate(speeds):
				lua.write("\t\t\t[{0}] = {1},\n".format(speed, i >= limit and "true" or "false").encode("utf-8"))

			lua.write(b"\t\t},\n")

		lua.write("\t}},\n}}\n\n-- Generated on {0} in {1:.3f} seconds\n".format(time.strftime("%a, %d %b %Y %H:%M:%SZ", time.gmtime()), time.clock() - start).encode("utf-8"))
except IOError:
	print("[ERROR] Could not write to CompanionData.lua")